import sys, os, shutil
from glob import glob
import random,math

def do_split(train_list, val_list, train_set_dir='train', val_set_dir='val'):
    work_dir = os.path.dirname(train_list[0])
    train_set_dir = os.path.join(work_dir, train_set_dir)
    val_set_dir = os.path.join(work_dir, val_set_dir)
    if os.path.exists(train_set_dir) is False:
        os.makedirs(train_set_dir)
    if os.path.exists(val_set_dir) is False:
        os.makedirs(val_set_dir)
    for file in train_list:
        shutil.move(file, os.path.join(train_set_dir, os.path.basename(file)))
        file = os.path.splitext(file)[0] + '.txt'
        shutil.move(file, os.path.join(train_set_dir, os.path.basename(file)))
    for file in val_list:
        shutil.move(file, os.path.join(val_set_dir, os.path.basename(file)))
        file = os.path.splitext(file)[0] + '.txt'
        shutil.move(file, os.path.join(val_set_dir, os.path.basename(file)))


def gen_pl_dataset_list(dir,rate=0.9):
    all_set = glob('{}/*jpg'.format(dir))
    random.shuffle(all_set)
    split_idx = math.ceil(len(all_set) * rate)
    train_set = all_set[:split_idx]
    val_set = all_set[split_idx:]
    do_split(train_set,val_set)

if __name__ == '__main__':
    # gen_pl_dataset_list('./anno_eg/')
    gen_pl_dataset_list('/home/leo/Downloads/datas/gen_wpod_det_trainset/mix')
